import torch
import matplotlib.pyplot as plt

fig = plt.figure()

x1 = torch.linspace(-10, 10, 100)
x2 = torch.linspace(-10, 10, 100)

x1, x2 = torch.meshgrid([x1, x2], indexing='ij')
print(x1.shape, x2.shape)


def sigmoid(x):
    return 1 / (1 + torch.exp(-x))


w1 = torch.normal(2, 0.2, size=(100, 100))
w2 = torch.normal(3, 0.2, size=(100, 100))
w3 = torch.normal(-1, 0.2, size=(100, 100))
w4 = torch.normal(-1, 0.2, size=(100, 100))
b = torch.tensor(2)

z1 = sigmoid(w1 * x1 + w2 * x2 + b)
z2 = sigmoid(w1 * x1 + w2 * x2 + b) + sigmoid(w3 * x1 + w4 * x2 + b)
ax1 = fig.add_subplot(121, projection="3d")
ax1.plot_surface(x1, x2, z1, cmap=plt.cm.YlGnBu_r)
ax2 = fig.add_subplot(122, projection="3d")
ax2.plot_surface(x1, x2, z2, cmap=plt.cm.YlGnBu_r)
plt.show()
